{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "extraordinary-confidentiality",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import glob\n",
    "from PIL import Image\n",
    "from torch.utils.data import Dataset\n",
    "import os\n",
    "import time\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import numpy as np\n",
    "import random\n",
    "import glob\n",
    "from skimage import io,transform,color\n",
    "from sklearn.utils import shuffle\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image\n",
    "import torch.optim as optim\n",
    "import torchvision.models as models\n",
    "from torch.utils.data import DataLoader\n",
    "from trash_dataloader import TrashDataset\n",
    "import time\n",
    "import seaborn\n",
    "import sys\n",
    "import glob\n",
    "import threading\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "from ev_toolkit import plot_tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "continuing-kennedy",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"class\": \"Harmful Waste_Battery board\"\n",
      "} {\n",
      "    \"class\": \"Harmful Waste_Battery board\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "from skimage import io,transform,color\n",
    "from torchvision import transforms, models\n",
    "# 自己的模型\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "MAX_EPOCH = 6\n",
    "BATCH_SIZE = 1\n",
    "LR = 0.002\n",
    "\n",
    "start_epoch = -1\n",
    "lr_decay_step = 7\n",
    "\n",
    "train_dir = '../../../../home/data/'\n",
    "num = os.listdir(train_dir)\n",
    "train_dir = os.path.join(train_dir,num[0])\n",
    "# 仅用于编码测试\n",
    "type_lst = os.listdir(train_dir)\n",
    "classes = type_lst\n",
    "norm_mean = [0.485, 0.456, 0.406]\n",
    "norm_std = [0.229, 0.224, 0.225]\n",
    "\n",
    "valid_transform = transforms.Compose([\n",
    "    transforms.Resize((112,112)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(norm_mean, norm_std),\n",
    "])\n",
    "# 构建MyDataset实例\n",
    "valid_data = TrashDataset(data_dir = train_dir,transform=valid_transform)\n",
    "# 构建DataLoder\n",
    "valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "\n",
    "def init():\n",
    "    # 测试时选择的文件名\n",
    "    pth = '../models/202158model_resnet_20epoch/models.pkl'\n",
    "    model = torch.load(pth)\n",
    "    model.to('cpu')\n",
    "    \n",
    "    return model\n",
    "\n",
    "# 根据训练的标签设置\n",
    "class_dict = {}\n",
    "f = open('./class.txt','r')\n",
    "a = f.read()\n",
    "class_dict = eval(a)\n",
    "class_dict = {value:key for key, value in class_dict.items()}\n",
    "f.close()\n",
    "\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision import transforms as T\n",
    "import torch as t\n",
    "\n",
    "\n",
    "\n",
    " # Resize：缩放\n",
    "\n",
    "\n",
    "\n",
    "def process_image(net, input_image, args=None):\n",
    "    img = input_image\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    img = Image.fromarray(img)\n",
    "    norm_mean = [0.485, 0.456, 0.406]\n",
    "    norm_std = [0.229, 0.224, 0.225]\n",
    "    transforms = T.Compose([T.Resize(112),T.ToTensor(),T.Normalize(norm_mean,norm_std)]) \n",
    "    img = transforms(img)\n",
    "    img = np.array(img)\n",
    "    img = img.transpose(0,2,1)\n",
    "    img = torch.tensor([img])\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        out = net(img)\n",
    "        _, pred = torch.max(out.data, 1)\n",
    "        data = json.dumps({'class': class_dict[pred[0].item()]},indent=4)\n",
    "    return data, out, pred\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    net = init()\n",
    "    labels = 139\n",
    "    path_img = '../../../../home/data/19/Harmful Waste_Battery board/bc2317b60aa246269f30d4b298344cb2.jpg'\n",
    "#     img = Image.open(path_img).convert('RGB') \n",
    "    img = cv2.imread(path_img)\n",
    "    inputs = np.asarray(img)\n",
    "#     inputs = torch.from_numpy(inputs)\n",
    "\n",
    "    dic, outputs, predicted = process_image(net, inputs)\n",
    "    print(dic, json.dumps({'class': class_dict[labels]},indent=4))\n",
    "    outputs = outputs.to('cpu')\n",
    "    predicted = predicted.to('cpu')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "precise-luxury",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = json.dumps({'class': class_dict[labels[0].item()]},indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "completed-auditor",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([110]) 0\n"
     ]
    }
   ],
   "source": [
    "print(predicted,labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "renewable-alpha",
   "metadata": {},
   "outputs": [],
   "source": [
    "type_lst = os.listdir(train_dir)\n",
    "classes = type_lst\n",
    "def load_data(type_lst,file_path):\n",
    "    n = 0\n",
    "    images=[] ##新建一个空列表用于存放图片数集\n",
    "    labels=[] ##新建一个空列表用于存放标签数集\n",
    "    trash_name = {}\n",
    "\n",
    "    for j in type_lst:\n",
    "        temp_path = os.path.join(file_path,j)\n",
    "        trash_name[j] = n\n",
    "        n += 1\n",
    "        # print('\\n' + '{} is finish!'.format(j))\n",
    "    return trash_name\n",
    "trash_name = load_data(type_lst,train_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "future-finding",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'KitchenWaste_mixed congee',\n",
       " 1: 'KitchenWaste_iceCream',\n",
       " 2: 'KitchenWaste_candiedGourdOnStick',\n",
       " 3: 'KitchenWaste_Coffee',\n",
       " 4: 'KitchenWaste_Cherry Tomatoes',\n",
       " 5: 'KitchenWaste_nut',\n",
       " 6: 'KitchenWaste_Chocolates',\n",
       " 7: 'KitchenWaste_jelly',\n",
       " 8: 'KitchenWaste_Walnut',\n",
       " 9: 'KitchenWaste_orange',\n",
       " 10: 'KitchenWaste_Leftovers',\n",
       " 11: 'KitchenWaste_Fruits',\n",
       " 12: 'KitchenWaste_pickled cabbage',\n",
       " 13: 'KitchenWaste_pitaya',\n",
       " 14: 'KitchenWaste_Roast Chicken',\n",
       " 15: 'KitchenWaste_melon seed',\n",
       " 16: 'KitchenWaste_tomato',\n",
       " 17: 'KitchenWaste_Straw cup',\n",
       " 18: 'KitchenWaste_Straw bowl',\n",
       " 19: 'KitchenWaste_Vermicelli',\n",
       " 20: 'KitchenWaste_strawberry',\n",
       " 21: 'KitchenWaste_pineapple',\n",
       " 22: 'KitchenWaste_French fries',\n",
       " 23: 'KitchenWaste_Potato chips',\n",
       " 24: 'KitchenWaste_Mushroom',\n",
       " 25: 'KitchenWaste_Egg Tart',\n",
       " 26: 'KitchenWaste_Bean curd',\n",
       " 27: 'KitchenWaste_Chicken wings',\n",
       " 28: 'Recyclable_Stainless Steel Products',\n",
       " 29: 'Recyclable_table tennis bat',\n",
       " 30: 'Recyclable_book',\n",
       " 31: 'Recyclable_Weighing scale',\n",
       " 32: 'Recyclable_vacuum cup',\n",
       " 33: 'Recyclable_envelope',\n",
       " 34: 'Recyclable_Charging head',\n",
       " 35: 'Recyclable_portable battery',\n",
       " 36: 'Recyclable_Rechargeable toothbrush',\n",
       " 37: 'Recyclable_Charging line',\n",
       " 38: 'Recyclable_Cycling',\n",
       " 39: 'Recyclable_card',\n",
       " 40: 'Recyclable_Desk lamp',\n",
       " 41: 'Recyclable_Tag',\n",
       " 42: 'Recyclable_hair drier',\n",
       " 43: 'Recyclable_globe',\n",
       " 44: 'Recyclable_Subway ticket',\n",
       " 45: 'Recyclable_solar heater',\n",
       " 46: 'Recyclable_ruler',\n",
       " 47: 'Recyclable_Nylon rope',\n",
       " 48: 'Recyclable_Hat',\n",
       " 49: 'Recyclable_Wrist watch',\n",
       " 50: 'Recyclable_Bracelet',\n",
       " 51: 'Recyclable_printer',\n",
       " 52: 'Recyclable_tyre pump',\n",
       " 53: 'Recyclable_Floor sweeping robot',\n",
       " 54: 'Recyclable_Empty bottle of skin care products',\n",
       " 55: 'Recyclable_Draw bar box',\n",
       " 56: 'Recyclable_slipper',\n",
       " 57: 'Recyclable_patch board',\n",
       " 58: 'Recyclable_radio',\n",
       " 59: 'Recyclable_hot pack',\n",
       " 60: 'Recyclable_telescope',\n",
       " 61: 'Recyclable_Wooden cutting board',\n",
       " 62: 'Recyclable_Cask',\n",
       " 63: 'Recyclable_Wooden comb',\n",
       " 64: 'Recyclable_Wooden spatula',\n",
       " 65: 'Recyclable_Wood carving',\n",
       " 66: 'Recyclable_pillow',\n",
       " 67: 'Recyclable_Table',\n",
       " 68: 'Recyclable_mould',\n",
       " 69: 'Recyclable_kettle',\n",
       " 70: 'Recyclable_sofa',\n",
       " 71: 'Recyclable_Fire Extinguisher',\n",
       " 72: 'Recyclable_Lampshade',\n",
       " 73: 'Recyclable_ashtray',\n",
       " 74: 'Recyclable_Hot water bottle',\n",
       " 75: 'Recyclable_Gas stove',\n",
       " 76: 'Recyclable_Gas cylinder',\n",
       " 77: 'Recyclable_Glassware',\n",
       " 78: 'Recyclable_Glass bottles',\n",
       " 79: 'Recyclable_Glass ball',\n",
       " 80: 'Recyclable_Yoga ball',\n",
       " 81: 'Recyclable_Electric shaver',\n",
       " 82: 'Recyclable_Electric curling stick',\n",
       " 83: 'Recyclable_Electronic scale',\n",
       " 84: 'Recyclable_Electric iron',\n",
       " 85: 'Recyclable_Electromagnetic furnace',\n",
       " 86: 'Recyclable_Television',\n",
       " 87: 'Recyclable_Circuit board',\n",
       " 88: 'Recyclable_electric fan',\n",
       " 89: 'Recyclable_rice cooker',\n",
       " 90: 'Recyclable_boarding pass',\n",
       " 91: 'Recyclable_plate',\n",
       " 92: 'Recyclable_bowl',\n",
       " 93: 'Recyclable_magnet',\n",
       " 94: 'Recyclable_humidifier',\n",
       " 95: 'Recyclable_cage',\n",
       " 96: 'Recyclable_playing cards',\n",
       " 97: 'Recyclable_network card',\n",
       " 98: 'Recyclable_headset',\n",
       " 99: 'Recyclable_Socks',\n",
       " 100: 'Recyclable_trousers',\n",
       " 101: 'Recyclable_Calculator',\n",
       " 102: 'Recyclable_Microphone',\n",
       " 103: 'Recyclable_Soybean Milk machine',\n",
       " 104: 'Recyclable_tyre',\n",
       " 105: 'Recyclable_Measuring cup',\n",
       " 106: 'Recyclable_Wire ball',\n",
       " 107: 'Recyclable_Aluminum products',\n",
       " 108: 'Recyclable_lid',\n",
       " 109: 'Recyclable_keyboard',\n",
       " 110: 'Recyclable_Tweezers',\n",
       " 111: 'Recyclable_Alarm',\n",
       " 112: 'Recyclable_Umbrella',\n",
       " 113: 'Recyclable_shoes',\n",
       " 114: 'Recyclable_sound',\n",
       " 115: 'Recyclable_Mat',\n",
       " 116: 'Recyclable_fish tank',\n",
       " 117: 'Recyclable_mouse',\n",
       " 118: 'Other Waste_PE plastic bag',\n",
       " 119: 'Other Waste_sanitary cup',\n",
       " 120: 'Other Waste_Disposable cotton swab',\n",
       " 121: 'Other Waste_Bamboo stick',\n",
       " 122: 'Other Waste_sticky note',\n",
       " 123: 'Other Waste_Kitchen gloves',\n",
       " 124: 'Other Waste_Kitchen cloth',\n",
       " 125: 'Other Waste_Mask',\n",
       " 126: 'Other Waste_Bath towel',\n",
       " 127: 'Other Waste_Shell',\n",
       " 128: 'Other Waste_towel',\n",
       " 129: 'Other Waste_Wet wipes',\n",
       " 130: 'Other Waste_Cleaning cloth',\n",
       " 131: 'Other Waste_adhesive tape',\n",
       " 132: 'Other Waste_Flyswatter',\n",
       " 133: 'Other Waste_Teapot pieces',\n",
       " 134: 'Other Waste_Straw hat',\n",
       " 135: 'Other Waste_Pregnancy Test Kit',\n",
       " 136: 'Other Waste_Feather duster',\n",
       " 137: 'Harmful Waste_Insecticide',\n",
       " 138: 'Harmful Waste_Dry battery',\n",
       " 139: 'Harmful Waste_Battery board',\n",
       " 140: 'Harmful Waste_The button battery',\n",
       " 141: 'Harmful Waste_glue',\n",
       " 142: 'Harmful Waste_Drug packaging',\n",
       " 143: 'Harmful Waste_pill',\n",
       " 144: 'Harmful Waste_ointment',\n",
       " 145: 'Harmful Waste_Battery'}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_dict"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
